import os
import numpy as np
import matplotlib.pyplot as plt

import sys

from numpy.lib.function_base import disp
sys.path.insert(0, '/media/tan/_dde_data/projects/ginkgo_vfife/scripts')
from visualization import ResultFileSystem, TimeHistory


def plot(data, figname):
    """ plot results of openVFIFE

    Args:
        data (ndarray): [t, x, y, z]
    """
    figsize = np.array([8, 6]) / 2.54
    fig, ax = plt.subplots(figsize=figsize, tight_layout=True)
    ax.plot(data[:,0], data[:,1], c="black", lw=1.0, label="Ux")
    ax.plot(data[:,0], data[:,2], c="red", lw=1.0, label="Uy")
    ax.legend(loc="upper right")
    plt.savefig(figname, dpi=300)
    plt.close(fig)


if __name__ == "__main__":

    # load data from openVFIFE result, differenct damping coeff
    dir = "/home/tan/Desktop/test/vfife_examples/example3/"

    fname = os.path.join(dir, "example3")
    file = ResultFileSystem(fname)
    th = TimeHistory(file)
    displace = th.extract_particle_motion(2, "displace")

    # save reults
    cwd = os.getcwd()
    figname = os.path.join(cwd, "displace.svg")
    plot(displace, figname)